import json
from collections import defaultdict
import numpy as np
import networkx as nx
import heapq
from flow1 import TrafficGenerator
from switchesalgo import allocate_switches_with_es
from topo1 import build_hybrid_topology_shared_tor


class HybridNetworkOptimizer:
    def __init__(self, network_topology, v_e=0.8, t_r=0.1):
        self.G = network_topology
        self.v_e = v_e  # 光电转换效率
        self.t_r = t_r  # 光链路重构时间
        self.optical_caps = self._get_optical_capacities()
        self.electrical_caps = self._get_electrical_capacities()

    def _get_optical_capacities(self):
        """获取所有光链路容量"""
        return {(u, v): data['capacity'] for u, v, data in self.G.edges(data=True)
                if data.get('type') == 'optical'}

    def _get_electrical_capacities(self):
        """获取所有电链路容量"""
        return {(u, v): data['capacity'] for u, v, data in self.G.edges(data=True)
                if data.get('type') == 'electrical'}

    def find_hybrid_path(self, source, target, flow, x, r):
        """
        混合路径查找算法（考虑光电协同）
        返回：(最优路径, 路径时延)
        """
        # 初始化优先队列 (累计时延, 当前节点, 路径, 是否使用光链路)
        heap = [(0, source, [source], False)]
        visited = {}
        best_delays = defaultdict(lambda: float('inf'))

        while heap:
            current_delay, u, path, optical_used = heapq.heappop(heap)

            if u == target:
                return path, current_delay

            if u in visited and visited[u] <= current_delay:
                continue
            visited[u] = current_delay

            for v in self.G.neighbors(u):
                edge_data = self.G[u][v]
                link_type = edge_data['type']
                new_optical = optical_used

                # 计算链路时延分量
                if link_type == 'optical':
                    C_o = self.optical_caps[(u, v)]
                    term1 = (flow * x[(u, v)] * r[(u, v)]) / C_o
                    term2 = ((1 - r[(u, v)]) * (flow - self.v_e * self.t_r) / C_o + self.t_r) * x[(u, v)]
                    link_delay = max(term1, term2)
                    new_optical = True
                else:
                    C_e = self.electrical_caps[(u, v)]
                    term3 = flow * (1 - x[(u, v)]) / C_e
                    link_delay = term3

                new_delay = max(current_delay, link_delay)

                if new_delay < best_delays[v]:
                    best_delays[v] = new_delay
                    heapq.heappush(heap, (new_delay, v, path + [v], new_optical))

        return None, float('inf')

    def optimize_reconfiguration(self, new_demand_matrix, old_paths, new_paths):
        """
        改进版动态流量重构优化
        返回: (重构路径, 路径时延, 非重构时延)
        """
        # 分类通信对
        recon_pairs = []
        no_recon_pairs = []
        for pair in new_demand_matrix:
            if old_paths.get(pair, []) != new_paths.get(pair, []):
                recon_pairs.append(pair)
            else:
                no_recon_pairs.append(pair)

        # 处理非重构通信对
        link_usage = defaultdict(int)
        for pair in no_recon_pairs:
            path = old_paths[pair]
            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]
                if self.G[u][v]['type'] == 'electrical':
                    link_usage[(u, v)] += new_demand_matrix[pair]

        # 批量更新电链路容量
        for (u, v), usage in link_usage.items():
            self.G[u][v]['capacity'] = max(self.G[u][v]['capacity'] - usage, 0)

        # 优化变量初始化 (示例实现，需结合实际优化器)
        x, r = self.initialize_optimization_vars(recon_pairs)

        # 处理重构通信对
        sorted_pairs = sorted(recon_pairs, key=lambda p: new_demand_matrix[p], reverse=True)
        recon_paths = {}
        delays = {}

        for pair in sorted_pairs:
            s, d = pair
            flow = new_demand_matrix[pair]
            path, delay = self.find_hybrid_path(s, d, flow, x, r)

            if path:
                # 更新网络状态
                self.update_network_capacity(path, flow, x, r)
                recon_paths[pair] = path
                delays[pair] = delay
            else:
                recon_paths[pair] = None
                delays[pair] = float('inf')

        # 计算非重构时延
        no_delays = self.calculate_non_reconfig_delays(no_recon_pairs, old_paths)

        return recon_paths, delays, no_delays

    def initialize_optimization_vars(self, recon_pairs):
        """优化变量初始化示例"""
        x = defaultdict(float)
        r = defaultdict(float)
        # TODO: 实际应通过求解器获得最优x和r
        # 这里示例随机初始化
        for u, v in self.G.edges():
            x[(u, v)] = 0.5 if self.G[u][v]['type'] == 'optical' else 0
            r[(u, v)] = 0.5
        return x, r

    def update_network_capacity(self, path, flow, x, r):
        """更新网络容量状态"""
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            if self.G[u][v]['type'] == 'optical':
                allocated = x[(u, v)] * r[(u, v)] * flow
                self.G[u][v]['C_o'] -= allocated
            else:
                self.G[u][v]['capacity'] -= flow * (1 - x[(u, v)])

    def calculate_non_reconfig_delays(self, pairs, old_paths):
        """计算非重构路径时延"""
        delays = {}
        for pair in pairs:
            path = old_paths[pair]
            max_delay = 0
            for i in range(len(path) - 1):
                u, v = path[i], path[i + 1]
                if self.G[u][v]['type'] == 'optical':
                    C_o = self.G[u][v]['C_o']
                    delay = self.G[u][v]['latency']  # 假设存在传播时延属性
                else:
                    C_e = self.G[u][v]['capacity']
                    flow = new_demand_matrix[pair]
                    delay = flow / C_e if C_e > 0 else float('inf')
                max_delay = max(max_delay, delay)
            delays[pair] = max_delay
        return delays


# ------------------- 使用示例 -------------------
if __name__ == "__main__":
    # 初始化网络拓扑
    G = build_hybrid_topology_shared_tor()

    # 创建优化器
    optimizer = HybridNetworkOptimizer(G, v_e=0.8, t_r=0.1)

    # 加载流量需求（示例）
    with open('traffic_simulation.json') as f:
        traffic_data = json.load(f)

    # TODO
    # 获取不同时间片的流量矩阵
    baseline_traffic = ...  # 基础流量
    peak_traffic = ...  # 突发流量

    # 计算基础路径
    S1, base_paths = allocate_switches_with_es(G, ...)

    # 计算峰值路径
    S2, peak_paths = allocate_switches_with_es(G, ...)

    # 执行优化
    recon_paths, delays, no_delays = optimizer.optimize_reconfiguration(
        peak_traffic, base_paths, peak_paths
    )

    # 输出结果
    print("重构路径时延分布：")
    for pair, delay in delays.items():
        print(f"{pair}: {delay:.2f} ms")

    print("\n链路剩余容量：")
    for u, v in G.edges():
        if G[u][v]['type'] == 'optical':
            print(f"光链路 {u}→{v}: {G[u][v]['C_o']} Gbps")
        else:
            print(f"电链路 {u}→{v}: {G[u][v]['capacity']} Gbps")